#![allow(dead_code)]

use anyhow::anyhow;
use axum::response::IntoResponse;
use chrono::{Datelike, Local};
use futures::StreamExt;
use http::header;
use regex::Regex;
use reqwest::Response;
use serde::{Deserialize, Serialize};
use sq_global::{GPT_API_HOST_CFG, GPT_API_KEY_CFG};
use std::sync::OnceLock;

static GPT_CLI: OnceLock<GptClient> = OnceLock::new();

pub fn init() {
    GPT_CLI
        .set(GptClient::init())
        .expect("GPT_CLI not initialized");
}

#[inline]
pub fn cli() -> &'static GptClient {
    GPT_CLI.get().expect("GPT_CLI not initialized")
}

#[derive(Debug, Serialize)]
struct ChatMessage {
    role: &'static str,
    content: String,
}

impl ChatMessage {
    fn with_context(context: String) -> Self {
        Self {
            role: "system",
            content: context,
        }
    }
    fn with_prompt(prompt: String) -> Self {
        Self {
            role: "user",
            content: prompt,
        }
    }
}

#[derive(Debug, Serialize, Default)]
pub struct ChatRequest {
    #[serde(default = "default_model")]
    model: &'static str,
    stream: bool,
    messages: Vec<ChatMessage>,
}

fn default_model() -> &'static str {
    "gpt-3.5-turbo"
}

impl ChatRequest {
    fn add(&mut self, msg: ChatMessage) {
        self.messages.push(msg);
    }
}

#[derive(Debug, Deserialize)]
struct ChatResponse {
    id: String,
    object: String,
    created: i64,
    model: String,
    choices: Vec<ChatChoice>,
}
#[derive(Debug, Deserialize)]
struct ChatChoice {
    index: u32,
    delta: ChatDelta,
    finish_reason: Option<String>,
}
#[derive(Debug, Deserialize)]
struct ChatDelta {
    role: Option<String>,
    content: Option<String>,
    refusal: Option<String>,
}

#[derive(Debug)]
pub struct GptClient {
    client: reqwest::Client,
    api_host: String,
    api_key: String,
    regex: Regex,
}

impl GptClient {
    pub fn init() -> GptClient {
        let api_host = sq_config::get_unchecked::<String>(GPT_API_HOST_CFG);
        let mut api_key = sq_config::get_unchecked::<String>(GPT_API_KEY_CFG);
        api_key.insert_str(0, "Bearer ");
        Self {
            client: sq_global::http_cli(),
            api_host,
            api_key,
            regex: Regex::new(r#""content":\s*"(.*?)""#).unwrap(),
        }
    }
    pub async fn complete(&self, req: &ChatRequest) -> anyhow::Result<Response> {
        let response = self
            .client
            .post(&self.api_host)
            .header(header::ACCEPT, "text/event-stream")
            .header(header::AUTHORIZATION, &self.api_key)
            .json(req)
            .send()
            .await?;

        // 确保响应状态成功
        if !response.status().is_success() {
            let error_text = response.text().await?;
            return Err(anyhow!("请求chatgpt失败！: {error_text}"));
        }
        Ok(response)
    }
    fn extract_content_from_json<'a>(&self, json_str: &'a str) -> Option<&'a str> {
        self.regex
            .captures(json_str)
            .and_then(move |caps| caps.get(1).map(|content| content.as_str()))
    }
}

pub async fn chat_stream(prompt: String) -> anyhow::Result<axum::response::Response> {
    let client = cli();
    let chat_req = prepare_req(prompt)?;
    let mut resp = String::new();
    let response = client.complete(&chat_req).await?;
    // 获取字节流
    let mut stream = response.bytes_stream();

    // 处理SSE流式响应
    while let Some(chunk_result) = stream.next().await {
        let chunk_bytes = chunk_result?;
        let text = String::from_utf8_lossy(&chunk_bytes);

        // 处理SSE格式的数据
        for line in text.lines() {
            if line.starts_with("data: ") {
                let data = line.trim_start_matches("data: ");
                // 跳过心跳消息
                if data == "[DONE]" {
                    continue;
                }

                let response: Option<ChatResponse> =
                    serde_json::from_str(data).ok().or_else(|| {
                        tracing::error!("获取到异常数据：{data}");
                        None
                    });

                let content = response
                    .as_ref()
                    .map(|r| &r.choices)
                    .and_then(|v| v.first())
                    .map(|c| &c.delta)
                    .and_then(|d| d.content.as_ref())
                    .filter(|&c| !c.is_empty())
                    .map(|v| v.as_str())
                    .or_else(|| {
                        client
                            .extract_content_from_json(data)
                            .filter(|&v| !v.is_empty())
                    });

                match content {
                    None => continue,
                    Some(slice) => {
                        resp.push_str(slice);
                    }
                }
            }
        }
    }
    Ok(resp.into_response())
}

fn prepare_req(prompt: String) -> anyhow::Result<ChatRequest> {
    let mut chat_request = ChatRequest::default();
    let curren_time = Local::now().naive_local();
    let date_str = curren_time.format("%Y-%m-%d %H:%M:%S");
    let number_from_monday = curren_time.weekday().number_from_monday();

    let weekday = match number_from_monday {
        1 => "周一",
        2 => "周二",
        3 => "周三",
        4 => "周四",
        5 => "周五",
        6 => "周六",
        7 => "周日",
        _ => return Err(anyhow!("计算周几异常！")),
    };
    let context = format!(
        "你是智能财务助手，注意要给到当前用户尊称。中国当前时间：{}，{}",
        date_str, weekday
    );
    chat_request.add(ChatMessage::with_context(context));
    chat_request.add(ChatMessage::with_prompt(prompt));
    Ok(chat_request)
}
